"""
模拟vggnet(串联结构），处理cifar10
标准的vggnet的输入是24*24，32位下采样之后，得到7*7的特征图
--> 调整
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class VGGBase(nn.Module):
    def __init__(self):
        super().__init__()
        # 输入：3*28*28 （因为自定义数据集的时候crop了将32裁剪到了28）
        self.conv1 = nn.Sequential(
            # 因为采用的核是3*3，不加padding输出的尺寸会发生变化
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )
        self.max_pooling1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 经过池化层后，尺寸变为14*14  。14 = （28-2(kenel)）/2(stride) + 1
        # 图像尺寸缩小，通常会将channel翻倍，
        self.conv2_1 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )

        self.conv2_2 = nn.Sequential(
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )
        self.max_pooling2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 此时输出的维度7*7 . 7= （14-2）/2 + 1

        # 7*7
        # 图像尺寸缩小，通常会将channel翻倍，
        self.conv3_1 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU()
        )

        self.conv3_2 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU()
        )
        self.max_pooling3 = nn.MaxPool2d(kernel_size=2, stride=2, padding=1)
        # 此时尺寸4*4(用padding补了一维） 4 = （7+2*1 - 2）/2 + 1

        # 输入：4*4
        self.conv4_1 = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU()
        )

        self.conv4_2 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU()
        )

        self.max_pooling4 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 此时尺寸为：2*2

        # 特征图的尺寸：batchsize*512*2*2 --> batchsize * (512 * 4)
        self.fc = nn.Linear(
            512 * 2 * 2,  # 通道数*图像的大小
            10,  # 10个类别
        )

    def forward(self, x):
        batchsize = x.size(0)  # 每次训练丢进来的图片数，在定义MyDataset的时候有设置
        out = self.conv1(x)
        out = self.max_pooling1(out)

        out = self.conv2_1(out)
        out = self.conv2_2(out)
        out = self.max_pooling2(out)

        out = self.conv3_1(out)
        out = self.conv3_2(out)
        out = self.max_pooling3(out)

        out = self.conv4_1(out)
        out = self.conv4_2(out)
        out = self.max_pooling4(out)

        out = out.view(batchsize, -1)
        # fc层需要将输入，转化为 batchsize * 每张图的向量化数字数量
        # batchsize * c * h * w --> batchsize * n
        out = self.fc(out)
        out = F.log_softmax(out, dim=1)

        return out


VGGNet = VGGBase
